import cv2
import numpy as np
import os
import matplotlib.pyplot as plt


def detect_frog_by_color(img_path, visualize=False):
    """
    基于颜色特征检测青蛙并返回ROI区域
    :param img_path: 图像文件路径
    :param visualize: 是否可视化处理过程
    :return: ROI区域图像和检测状态
    """
    # 读取图像
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图像: {img_path}")
        return None, False

    # 转换为HSV颜色空间（对光照变化更鲁棒）
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # 增强图像对比度（改善颜色检测）
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    hsv[:, :, 2] = clahe.apply(hsv[:, :, 2])

    # 定义青蛙颜色的HSV范围（绿色、暗黄色）
    # 扩大绿色范围：H:25-95 (覆盖黄绿到青绿)
    lower_green = np.array([25, 40, 40])
    upper_green = np.array([95, 255, 255])

    # 调整黄色范围更暗：降低V值上限 (使黄色更暗)
    lower_dark_yellow = np.array([10, 60, 30])  # 降低V值下限
    upper_dark_yellow = np.array([40, 255, 180])  # 降低V值上限

    # 创建颜色掩膜
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_dark_yellow = cv2.inRange(hsv, lower_dark_yellow, upper_dark_yellow)

    # 合并所有掩膜
    combined_mask = cv2.bitwise_or(mask_green, mask_dark_yellow)

    # 形态学操作（开运算去除小噪声，闭运算连接相邻区域）
    kernel = np.ones((7, 7), np.uint8)  # 增大内核尺寸以更好处理噪声
    cleaned_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
    cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(cleaned_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 如果没有找到轮廓，返回失败
    if not contours:
        return None, False

    # 找到最大轮廓（假设青蛙是最大的有色物体）
    largest_contour = max(contours, key=cv2.contourArea)

    # 获取边界框
    x, y, w, h = cv2.boundingRect(largest_contour)

    # 确保边界框合理大小（至少占图像的5%）
    min_size = 0.05 * img.shape[0] * img.shape[1]
    if w * h < min_size:
        return None, False

    # 扩展边界框（增加15%的边距）
    margin = 0.15
    x = max(0, int(x - margin * w))
    y = max(0, int(y - margin * h))
    w = min(img.shape[1] - x, int(w * (1 + 2 * margin)))
    h = min(img.shape[0] - y, int(h * (1 + 2 * margin)))

    # 提取ROI
    roi = img[y:y + h, x:x + w]

    # 可视化处理过程
    if visualize:
        # 创建可视化图像
        vis_img = img.copy()
        cv2.rectangle(vis_img, (x, y), (x + w, y + h), (0, 255, 0), 3)

        # 创建彩色掩膜显示
        color_mask = np.zeros_like(img)
        color_mask[mask_green > 0] = [0, 255, 0]  # 绿色区域
        color_mask[mask_dark_yellow > 0] = [0, 165, 255]  # 橙色表示暗黄色区域

        # 组合显示
        top_row = np.hstack((img, cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)))
        bottom_row = np.hstack((color_mask, vis_img))
        combined_vis = np.vstack((top_row, bottom_row))

        # 保存可视化结果
        output_dir = os.path.dirname(img_path) + "_visualization"
        os.makedirs(output_dir, exist_ok=True)
        vis_path = os.path.join(output_dir, f"vis_{os.path.basename(img_path)}")
        cv2.imwrite(vis_path, combined_vis)
        print(f"已保存可视化结果: {vis_path}")

    return roi, True


def preprocess_images_with_color(image_dir, output_dir):
    """
    使用颜色分割预处理图像：只保存成功检测到青蛙的ROI
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 创建日志文件
    log_file = os.path.join(output_dir, "processing_log.txt")
    with open(log_file, "w") as f:
        f.write("图像处理日志\n")
        f.write("=" * 50 + "\n")

    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    saved_count = 0

    print(f"开始处理 {len(image_files)} 张图像...")

    for img_file in image_files:
        img_path = os.path.join(image_dir, img_file)
        roi, detected = detect_frog_by_color(img_path, visualize=True)

        if detected and roi is not None:
            roi_path = os.path.join(output_dir, img_file)
            cv2.imwrite(roi_path, roi)
            saved_count += 1
            print(f"已保存ROI: {roi_path}")
            with open(log_file, "a") as f:
                f.write(f"成功: {img_file}\n")
        else:
            print(f"未检测到青蛙: {img_file}")
            with open(log_file, "a") as f:
                f.write(f"失败: {img_file}\n")

    success_rate = saved_count / len(image_files) * 100

    print(f"\n处理完成: 共处理 {len(image_files)} 张图片")
    print(f"成功检测并保存 {saved_count} 张青蛙ROI图像")
    print(f"未能检测到青蛙 {len(image_files) - saved_count} 张")
    print(f"成功率: {success_rate:.2f}%")

    with open(log_file, "a") as f:
        f.write("\n" + "=" * 50 + "\n")
        f.write(f"总图像数: {len(image_files)}\n")
        f.write(f"成功检测数: {saved_count}\n")
        f.write(f"失败检测数: {len(image_files) - saved_count}\n")
        f.write(f"成功率: {success_rate:.2f}%\n")

    return success_rate


if __name__ == '__main__':
    image_directory = 'D:/Frog/image3'
    output_directory = 'D:/Frog/BigFrog'

    # 处理图像
    success_rate = preprocess_images_with_color(image_directory, output_directory)
    print(f"颜色分割检测完成！成功率: {success_rate:.2f}%")